import socket
import time
import threading
from collections import deque

# 配置参数
WINDOW_SIZE = 100  # 滑动窗口大小
NACK_TIMEOUT = 0.5  # NACK超时时间（秒）
MAX_RETRIES = 3  # 最大重传次数


class NACKSender:
    def __init__(self, host, port):
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.address = (host, port)
        self.seq_num = 0
        self.send_window = deque()  # 发送窗口：[(seq_num, data, timestamp, retries), ...]
        self.nack_queue = deque()  # NACK请求队列
        self.running = False
        self.thread = None

    def start(self):
        """启动NACK处理线程"""
        self.running = True
        self.thread = threading.Thread(target=self._nack_processor)
        self.thread.daemon = True
        self.thread.start()

    def stop(self):
        """停止NACK处理线程"""
        self.running = False
        if self.thread:
            self.thread.join()

    def send_data(self, data):
        """发送数据并添加到发送窗口"""
        # 等待窗口有空间（简化实现，实际可能需要更复杂的流控）
        while len(self.send_window) >= WINDOW_SIZE:
            time.sleep(0.01)

        self.seq_num = (self.seq_num + 1) % 65536  # 防止序列号溢出
        packet = (self.seq_num, data)
        self._send_packet(packet)
        self.send_window.append((self.seq_num, data, time.time(), 0))

    def _send_packet(self, packet):
        """实际发送数据包（添加序列号头部）"""
        seq_num, data = packet
        header = seq_num.to_bytes(2, 'big')  # 2字节序列号
        self.socket.sendto(header + data, self.address)

    def _nack_processor(self):
        """处理NACK请求和超时重传的线程"""
        while self.running:
            # 处理NACK请求
            while self.nack_queue:
                missing_seq = self.nack_queue.popleft()
                self._retransmit(missing_seq)

            # 处理超时未确认的包
            current_time = time.time()
            for i in range(len(self.send_window)):
                seq, data, timestamp, retries = self.send_window[i]
                if current_time - timestamp > NACK_TIMEOUT and retries < MAX_RETRIES:
                    self.send_window[i] = (seq, data, current_time, retries + 1)
                    self._send_packet((seq, data))

            time.sleep(0.01)  # 避免CPU空转

    def handle_nack(self, nack_data):
        """处理接收到的NACK消息"""
        # 假设NACK消息格式为：前2字节为NACK数量，后续每2字节为一个缺失的序列号
        nack_count = int.from_bytes(nack_data[:2], 'big')
        for i in range(nack_count):
            missing_seq = int.from_bytes(nack_data[2 + 2 * i:4 + 2 * i], 'big')
            self.nack_queue.append(missing_seq)

    def _retransmit(self, seq_num):
        """重传指定序列号的数据包"""
        for i, (seq, data, _, retries) in enumerate(self.send_window):
            if seq == seq_num and retries < MAX_RETRIES:
                self.send_window[i] = (seq, data, time.time(), retries + 1)
                self._send_packet((seq, data))
                break


class NACKReceiver:
    def __init__(self, host, port):
        self.socket = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
        self.socket.bind((host, port))
        self.expected_seq = 0  # 期望接收的下一个序列号
        self.received_packets = {}  # 已接收但乱序的数据包：{seq_num: data}
        self.nack_timer = None
        self.running = False
        self.thread = None

    def start(self):
        """启动接收线程"""
        self.running = True
        self.thread = threading.Thread(target=self._receive_loop)
        self.thread.daemon = True
        self.thread.start()

    def stop(self):
        """停止接收线程"""
        self.running = False
        if self.thread:
            self.thread.join()

    def _receive_loop(self):
        """接收数据包的主循环"""
        while self.running:
            try:
                data, addr = self.socket.recvfrom(65535)
                if len(data) < 2:  # 至少包含2字节序列号
                    continue

                seq_num = int.from_bytes(data[:2], 'big')
                payload = data[2:]

                self._handle_packet(seq_num, payload, addr)
            except Exception as e:
                print(f"接收错误: {e}")
                time.sleep(0.1)

    def _handle_packet(self, seq_num, payload, sender_addr):
        """处理接收到的数据包"""
        # 如果是期望的序列号，直接处理
        if seq_num == self.expected_seq:
            self._process_data(payload)
            self.expected_seq = (self.expected_seq + 1) % 65536

            # 检查已接收的乱序包，看是否有可以连续处理的
            while self.expected_seq in self.received_packets:
                data = self.received_packets.pop(self.expected_seq)
                self._process_data(data)
                self.expected_seq = (self.expected_seq + 1) % 65536

            # 发送NACK请求缺失的包
            self._send_nack(sender_addr)

        # 如果是乱序包，保存到缓存
        elif (seq_num > self.expected_seq) or (
                seq_num < self.expected_seq and seq_num < 1000 and self.expected_seq > 64000):
            self.received_packets[seq_num] = payload

            # 定期发送NACK请求缺失的包
            if not self.nack_timer or time.time() - self.nack_timer > NACK_TIMEOUT:
                self._send_nack(sender_addr)
                self.nack_timer = time.time()

    def _send_nack(self, sender_addr):
        """发送NACK请求缺失的数据包"""
        # 收集所有缺失的序列号
        missing_seqs = []
        current = self.expected_seq

        # 查找窗口范围内的缺失包
        for _ in range(WINDOW_SIZE):
            if current not in self.received_packets:
                missing_seqs.append(current)
            current = (current + 1) % 65536
            if current == self.expected_seq:  # 避免无限循环
                break

        # 构造NACK消息
        if missing_seqs:
            nack_data = len(missing_seqs).to_bytes(2, 'big')
            for seq in missing_seqs:
                nack_data += seq.to_bytes(2, 'big')

            self.socket.sendto(nack_data, sender_addr)

    def _process_data(self, data):
        """处理完整接收到的数据（由子类实现）"""
        print(f"处理数据: {len(data)} 字节")


# 使用示例
if __name__ == "__main__":
    # 发送端示例
    sender = NACKSender("127.0.0.1", 9999)
    sender.start()

    # 接收端示例
    receiver = NACKReceiver("127.0.0.1", 10000)
    receiver.start()

    # 模拟发送数据
    try:
        for i in range(1000):
            sender.send_data(f"数据包 {i}".encode())
            time.sleep(0.01)
    except KeyboardInterrupt:
        sender.stop()
        receiver.stop()